# -*- coding: utf-8 -*-
# @Time    : 2025/3/17 15:41
# @Author  : sjh
# @Site    :
# @File    : main.py
# @Comment :
import json

import cv2
import numpy as np
import open3d as o3d

from utils.get_img import DepthReceiver

image_width = 640
depth_global = None
show_ply = False  # 是否显示点云


def depth2xyz(depth_map, depth_cam_matrix, flatten=False, depth_scale=1000):
    """ 将深度图转换为点云 """
    fx, fy = depth_cam_matrix[0, 0], depth_cam_matrix[1, 1]
    cx, cy = depth_cam_matrix[0, 2], depth_cam_matrix[1, 2]

    # 生成网格
    h, w = np.mgrid[0:depth_map.shape[0], 0:depth_map.shape[1]]

    # 计算 3D 坐标
    z = depth_map / depth_scale  # 归一化深度（单位: m）
    x = (w - cx) * z / fx
    y = (h - cy) * z / fy
    xyz = np.dstack((x, y, z)) if flatten == False else np.dstack((x, y, z)).reshape(-1, 3)

    return xyz


def xy_3d(x, y, depth, fx=238.494, fy=233.194, cx=229.981, cy=151.422, min_depth=0.1, max_depth=10.0, max_step=5):
    """
    计算 3D 坐标 (X, Y, Z)
    :param x: 像素 x 坐标
    :param y: 像素 y 坐标
    :param depth: 深度图 (单位: 米)
    :param fx: 焦距 fx (像素)
    :param fy: 焦距 fy (像素)
    :param cx: 光心 x 坐标
    :param cy: 光心 y 坐标
    :param baseline: 双目相机基线长度 (单位: mm)
    :param min_depth: 允许的最小深度值
    :param max_depth: 允许的最大深度值
    :param max_step: 搜索邻域最大步长
    :return: 3D 坐标 [X, Y, Z]
    """
    height, width = depth.shape
    step = 1

    # 获取当前点的深度
    Z = depth[y, x]

    # 如果深度无效，则尝试寻找邻域有效最小深度
    while (Z < min_depth or Z > max_depth) and step < max_step:
        surrounding_values = depth[max(0, y - step):min(height, y + 1 + step),
                             max(0, x - step):min(width, x + 1 + step)]
        valid_values = surrounding_values[(surrounding_values > min_depth) & (surrounding_values < max_depth)]
        if valid_values.size > 0:
            Z = np.min(valid_values)  # 选取最小有效深度
            break
        step += 1

    # 如果仍然无效，则返回 NaN
    if Z < min_depth or Z > max_depth:
        return np.array([np.nan, np.nan, np.nan])

    # 计算 3D 坐标
    X = (x - cx) * Z / fx
    Y = (y - cy) * Z / fy

    return np.array([X, Y, Z])


def on_mouse(event, x, y, flags, param):
    global depth_global
    if event == cv2.EVENT_LBUTTONDOWN:
        x = x if x <= image_width else x - image_width
        points_3d = xy_3d(x, y, depth_global)
        points_3d = points_3d.reshape(3, 1)
        print('点 (%d, %d) 的三维坐标 (%f, %f, %f)' % (
            x, y, points_3d[0], points_3d[1], points_3d[2]))



if __name__ == '__main__':
    receiver = DepthReceiver()
    keypoints_3d_list = []
    while True:
        bgr_image, depth = receiver.receive_frame()
        if bgr_image is None and depth is None:
            # print("Connection lost.")
            continue  # 退出循环
        depth_global = depth.astype(np.float32) / 1000.0
        depth_colormap = receiver.display_depth(depth)

        combined_image = np.hstack((bgr_image, depth_colormap))
        cv2.imshow("Estimated Image", combined_image)
        cv2.setMouseCallback("Estimated Image", on_mouse)

        if cv2.waitKey(1) & 0xFF == 27:  # 按 ESC 退出
            break


        # **显示点云并允许选点**
        if show_ply:
            depth_cam_matrix = np.array([[238.494, 0, 229.981],
                                         [0, 233.194, 151.422],
                                         [0, 0, 1]])
            max_depth = 10  # 单位 m
            depth_scale = 1000  # 假设深度图以 mm 存储
            depth_map = depth
            depth_map = np.where(depth_map / depth_scale > max_depth, 0, depth_map)
            depth_map = np.where(depth_map / depth_scale < 0, 0, depth_map)
            pc = depth2xyz(depth_map, depth_cam_matrix, flatten=True)

            # **可视化点云并允许选点**
            pcd = o3d.geometry.PointCloud()
            pcd.points = o3d.utility.Vector3dVector(pc)
            vis = o3d.visualization.VisualizerWithEditing()
            vis.create_window(window_name="Click to Get Depth")
            vis.add_geometry(pcd)

            print("\n请在窗口中 **按住Shift + 左键** 选点，然后关闭窗口后查看选中的点。\n")

            vis.run()  # 运行可视化（允许选点）
            vis.destroy_window()

            # **获取选中的点**
            picked_points = vis.get_picked_points()
            if picked_points:
                print("\n选中的 3D 点坐标：")
                for i, idx in enumerate(picked_points):
                    x, y, z = pc[idx]
                    print(f"点 {i + 1}: X={x:.3f}, Y={y:.3f}, Z={z:.3f} m")
            else:
                print("未选中任何点。")
    receiver.close()
    cv2.destroyAllWindows()
